from IdentityBasedEncryption import *
from functools import reduce


class Waters(IdentityBasedEncryption):
    def __init__(self, lm = 16):
        super(Waters, self).__init__()
        if lm not in [16, 32]:
            raise ValueError('only support 128 bits (16 bytes) or 256 bits (32 bytes) message')
        else:
            self.__lm = lm
            self.__hash = md5 if lm == 16 else sha256

    def __h1(self, b):
        """ {0, 1} ^ * -> {0, 1} ^ n """
        return bytes.fromhex(self.__hash(b).hexdigest())

    def system_build(self):
        s = self.group.random(ZR)
        P1, P2, u_minute = s * self.g, self.group.random(G1), self.group.random(G1)
        U = [self.group.random(G1) for _ in range(self.__lm * 8)]
        return {'s': s,
                'P1': P1,
                'P2': P2,
                'u_minute': u_minute,
                'U': U}

    def key_extract(self, user_id, sys_params):
        U = []
        for i, c in enumerate(bytes2bits(user_id)):
            if c == '1':
                U.append(i)
        r = self.group.random(ZR)
        sum_ui = reduce(lambda ui, uj: ui + uj, [sys_params['U'][i] for i in U])
        SU = {'S1': sys_params['s'] * sys_params['P2'] + r * (sys_params['u_minute'] + sum_ui),
              'S2': r * self.g,
              'user_id': user_id,
              'sum_ui': sum_ui}

        # randomize user secret key if necessary
        r_minute = self.group.random(ZR)
        SU_minute = {'S1': sys_params['s'] * sys_params['P2'] + (r + r_minute) * (sys_params['u_minute'] + sum_ui),
                     'S2': (r + r_minute) * self.g,
                     'user_id': user_id,
                     'sum_ui': sum_ui}

        # SU_minute is SU after re-randomization, that mean return SU is also a correct result
        return SU_minute  # or SU

    def encryption(self, m, sys_params, user_params):
        t = self.group.random(ZR)
        return m * (e(sys_params['P1'], sys_params['P2']) ** t), t * self.g, t * (
                    sys_params['u_minute'] + user_params['sum_ui'])

    def decryption(self, ciphertext, user_params):
        C1, C2, C3 = ciphertext
        S1, S2 = user_params['S1'], user_params['S2']
        return C1 * (e(S2, C3) / e(S1, C2))

if __name__ == '__main__':
    waters = Waters(16)
    waters.test(waters.group.random(GT), b'xxx')
